project(mnist)

# Torch
find_package(Torch REQUIRED PATHS $ENV{Torch_DIR} HINTS
             $ENV{HOME}/libtorch/share/cmake/Torch)
# Download the MNIST dataset
find_package(
  Python3
  COMPONENTS Interpreter
  REQUIRED)
option(DOWNLOAD_MNIST "Download the MNIST dataset from the internet" ON)
if(DOWNLOAD_MNIST)
  message(STATUS "Downloading MNIST dataset")
  # cmake-format: off
               execute_process(
                 COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tools/download_mnist.py 
                         -d ${CMAKE_BINARY_DIR}/data 
                         ERROR_VARIABLE 
                         DOWNLOAD_ERROR)
               # cmake-format: on
  if(DOWNLOAD_ERROR)
    message(FATAL_ERROR "Error downloading MNIST dataset: ${DOWNLOAD_ERROR}")
  endif()
endif()

# MPI_Cifar10
create_target(
  test_mpi_bgv_Cifar10_base
  "test_mpi_bgv_Cifar10_base.cpp;mpi_Cifar10_base.cpp;mpi_bgv_Cifar10_base.cpp;${CMAKE_SOURCE_DIR}/test/CNN/cifar10.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
  "helib;mpi;mkl")
target_include_directories(test_mpi_bgv_Cifar10_base PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_features(test_mpi_bgv_Cifar10_base PUBLIC cxx_range_for)
target_link_libraries(test_mpi_bgv_Cifar10_base ${TORCH_LIBRARIES})

create_target(
  test_mpi_bgv_Cifar10_test
  "test_mpi_bgv_Cifar10_test.cpp;mpi_Cifar10_test.cpp;mpi_bgv_Cifar10_test.cpp;${CMAKE_SOURCE_DIR}/test/CNN/cifar10.cpp;${MPI_SRC_PATH};${MPI_SRC_PATH}/mpi_util.cpp;${GENERAL_SRCS};${HYPERCUBE_SRCS}"
  "helib;mpi;mkl")
target_include_directories(test_mpi_bgv_Cifar10_test PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_features(test_mpi_bgv_Cifar10_test PUBLIC cxx_range_for)
target_link_libraries(test_mpi_bgv_Cifar10_test ${TORCH_LIBRARIES})

# create_target(train_cifar10 "cifar10_train.cpp;cifar10.cpp" "")
# target_include_directories(train_cifar10 PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(train_cifar10 PUBLIC cxx_range_for)
# target_link_libraries(train_cifar10 ${TORCH_LIBRARIES})

# create_target(bgv_cifar10 "bgv_cifar10.cpp;cifar10.cpp" "")
# target_include_directories(bgv_cifar10 PRIVATE ${TORCH_INCLUDE_DIRS})
# target_compile_features(bgv_cifar10 PUBLIC cxx_range_for)
# target_link_libraries(bgv_cifar10 ${TORCH_LIBRARIES})

if(MSVC)
  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
  add_custom_command(
    TARGET mnist
    POST_BUILD
    COMMAND ${CMAKE_COMMAND} -E copy_if_different ${TORCH_DLLS}
            $<TARGET_FILE_DIR:mnist>)
endif(MSVC)
